import concurrent.futures
import time
import random

import numpy as np
import argparse


def stream_requests(ques, output_file):
    pass


def measure_latency(ques, output_file, is_stream=False):
    start_time = time.time()
    _ = list(stream_requests(ques, output_file))
    end_time = time.time()
    return end_time - start_time


def perform_load_test(concurrency, total_requests, questions, output_file, is_stream=False):
    latencies = []
    questions = ["什么是三大专项", "江苏高三物生地，军校能不能报，哪些专业不能报", "山东文科在江苏怎么选学校",
                 "东南大学化学工程与工艺，生物科学，制药工程分流哪个好？", "男生高三物化地，辽宁，学日语好选学校吗"]
    with concurrent.futures.ThreadPoolExecutor(max_workers=concurrency) as executor:
        future_to_request = {executor.submit(measure_latency, random.choice(questions), output_file, is_stream): i for i
                             in range(total_requests)}
        for future in concurrent.futures.as_completed(future_to_request):
            try:
                latency = future.result()
                latencies.append(latency)
            except Exception as e:
                print(f"请求执行异常: {e}")

    # 计算统计数据
    p99 = np.percentile(latencies, 99)
    p95 = np.percentile(latencies, 95)
    total_time = sum(latencies)
    qps = total_requests / total_time

    return latencies, p99, p95, qps


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='I can do anything.')
    args = parser.parse_args()
    latencies, p99, p95, qps = perform_load_test(args.concurrency, args.total_requests, questions, output_file,
                                                 args.stream)

    # 打印统计结果
    print(f"延迟P99: {p99} 秒")
    print(f"延迟P95: {p95} 秒")
    print(f"QPS: {qps} 请求/秒")
